from ..builder import build_transform


class Compose:
    """
    数据增强组合
    """
    def __init__(self, transforms: list, tasks):
        assert isinstance(transforms, list)#[数据增强模块参数]
        self.transforms = []
        for transform in transforms:
            assert isinstance(transform, dict)
            transform = build_transform(transform)
            # 检查该数据增强是否在当前任务中
            fit_task = True
            for aug_task in transform.tasks:
                if aug_task not in tasks:
                    fit_task = False
                    break
            if fit_task is True:
                self.transforms.append(transform)
        # 按照优先级排序
        self.transforms.sort()

    def __call__(self, data):
        for t in self.transforms:
            data = t(data)
            if data is None:
                return None
        return data

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            format_string += '\n'
            format_string += f'    {t}'
        format_string += '\n)'
        return format_string
